﻿# coding: utf-8

# Copyright (c) P SOFTHOUSE. All rights reserved.

from collections import defaultdict

import maya.cmds as cmds

# ----------------------------------------------------------------------------
# Use this class to store and retrieve data for copy-paste operations.
# Also handles registering paste buttons to enable them when copy
# data is stored.
class CopyDataBuffer():

    # defaultdict where the key is the data buffer name (string)
    # and value is a list of associated Paste button names (strings)
    _pasteButtonNames = defaultdict(set)
    
    # defaultdict where the key is the data buffer name (string)
    # and the value is a string containing the copy buffer data.
    _dataBuffer = dict()
    
    # Clears all the class data
    @classmethod
    def clearAll(cls):
        cls._pasteButtonNames = defaultdict(set)
        cls._dataBuffer = dict()
    
    # Check if the specified copy data buffer contains copy data.
    # 
    # dataBufName (string): Name of the copy data buffer.
    # return: True if data was registered for the buffer. False otherwise.
    @classmethod
    def hasCopyData(cls, dataBufName):
        return dataBufName in cls._dataBuffer
    
    # Register a paste button to be associated to the specified 
    # copy data buffer.
    # The button will be enabled when data is registered for the buffer
    # (or enabled right away if there already is data available).
    # The button will be disabled if there is no data available.
    #
    # dataBufName (string): Name of the copy data buffer.
    # pasteButtonName (string): Name of the paste button in Maya's UI.
    @classmethod
    def registerPasteButton(cls, dataBufName, pasteButtonName):
        # If we already have copy data for this, we can simply enable the
        # button and forget about it.
        if dataBufName in cls._dataBuffer:
            cmds.button(pasteButtonName, edit=True, enable=True)
            
        else:
            # We are adding to a set, so duplicates will be ignored.
            cls._pasteButtonNames[dataBufName].add(pasteButtonName)
            
            # It is possible that the button is enabled even though we don't
            # have copy data for it in the case where we created the button,
            # copied some data (enabling the button), unloaded Pencil+ 
            # (clearing the copy data), reloaded Pencil+ and showed
            # the button again. In that case, the button created before
            # reloading Pencil+ is re-used so it will still be enabled
            # and we need to disable it.
            cmds.button(pasteButtonName, edit=True, enable=False)
        
    # Register copy data for the specified copy data buffer, overwriting 
    # existing data.
    # Also automatically activates the Paste buttons registered to the 
    # copy data buffer.
    #
    # dataBufName (string): Name of the copy data buffer.
    # data (string): Copy data string. This string will be returned as is
    #  when calling getCopyData for the same copy data buffer name. So the
    #  caller is responsible for managing the data format.
    @classmethod
    def setCopyData(cls, dataBufName, data):
        # Store the data
        cls._dataBuffer[dataBufName] = data
        
        # Activate the paste buttons
        if dataBufName in cls._pasteButtonNames:
            # We can forget about activated buttons.
            buttons = cls._pasteButtonNames.pop(dataBufName)
            for b in buttons:
                # Check if the button still exists
                if cmds.button(b, exists=True):
                    cmds.button(b, edit=True, enable=True)
    
    # Returns the copy data for the specified data buffer.
    #
    # dataBufName (string): Name of the copy data buffer.
    # Return (string): Copy data string for the specified data buffer. 
    #  This string will be the same as the one registered with setCopyData.
    #  If there is an error, an error message will be displayed and an empty
    #  string ("") will be returned.
    @classmethod
    def getCopyData(cls, dataBufName):
        # Check that we have data for this buffer name
        if dataBufName not in cls._dataBuffer:
            cmds.error("Could not find copy buffer data for buffer name=" + dataBufName)
            return ""
            
        return cls._dataBuffer[dataBufName]